%cd('/Users/tew417/Dropbox (QMUL-SEF)/POPI/POPIcodes_sigma2')
clc;clear all;close all;
global alpha betta depr sigmaC sigmaL phi govexp omega taumax N T ALPHA popweight
global kstst lstst cstst rstst wstst kstart kistart tauK0
global Delta1 Delta2 lamda PVtaxes deficit kg lam kmax

%run mainbench.m;

load('benchN2.mat')

kistart = kj;
kstart = sum(popweight.*kistart)+kg;
kmax=3;

%ALPHA=1;

% if (N==5)
% Delta1 = .5;
% Delta2 = [0.3 0.1 -0.1 -0.3];
% ALPHA=[0.9 0.8 0.7 0.6]
% end

lamda=lam;
Delta1 = .3;
Delta2 = -0.1;

kstst = fzero('findstst',kstart+0.3,options)

rstst = betta^(-1)-1+depr; %r from Euler at the steady state
estst = (rstst/alpha)^(1/(1-alpha))*kstst; % express total effective labor from r=F_k
for jj=2:N
ljpart(jj-1)=popweight(jj)*phi(jj)*lam(jj-1)^(-sigmaC/sigmaL)*(phi(jj)/phi(1))^(1/sigmaL);
end
lstst=estst/(popweight(1)*phi(1)+sum(ljpart)); %express l1 using the expression for total effective labor and l2=f(lam,l1)
for jj=2:N
    L2a(jj-1) = lamda(jj-1)^(-sigmaC/sigmaL)*(phi(jj)/phi(1))^(1/sigmaL); %last terms in (11) and f_l
end

T=150; %we assume that the economy reaches the steady state after T periods

tauK0 = taumax;

%initial guesses 
for i=1:T;
    time1 = i/T;
    X(i) = time1*(kstst) + (1-time1)*kstart; %k
    X(T+i) = lstst; %l
    if (N==5)
    X(2*T+i) = max(2-.07*i,0); %gamma
    end
    if (N==2)
    X(2*T+i) = max(1-.05*i,0); %gamma
    end
end;
X(3*T+1:3*T+1+2*(N-1)) = [Delta1 Delta2 lamda];

ALPHall=[0.4861 0.48 0.475 0.473 0.472 0.46 0.455 0.453 0.452 0.451 0.44 0.435 0.433 0.432 0.418 0.417 0.416 0.415 0.414 0.4 0.399 0.398 0.385 0.383 0.382 0.37 0.369 0.368 0.356 0.355 0.347 0.3467]

for ii=1:length(ALPHall)

    ALPHA=ALPHall(ii)
    
    if (ii>1)
    X=XX(ii-1,:);
    if ii==8 | ii==12 | ii==15 | ii==16 | ii==18 | ii==22 | ii==23 | ii==32
        X=XX(ii-2,:);
    end
    if ii==21
        X=XX(ii-3,:);
    end
    if ii==27
        X=XX(ii-4,:);
    end
    if ii==28 | ii==29
        X=0.5*XX(ii-1,:)+0.5*XX(ii-4,:);
    end
    if ii==30 | ii==31
        X=0.5*XX(ii-3,:)+0.5*XX(ii-6,:);
    end
    Delta1 = X(3*T+1);
    Delta2 = X(3*T+2:3*T+2+N-2);
    lamda = X(3*T+3+N-2:3*T+3+2*(N-2));
    end
    for i=1:T;
        time1 = i/T;
        X(i) = time1*(kstst) + (1-time1)*kstart; %k
    end
kstst = fzero('findstst',kstart+0.3)
rstst = betta^(-1)-1+depr; %r from Euler at the steady state
estst = (rstst/alpha)^(1/(1-alpha))*kstst; % express total effective labor from r=F_k
for jj=2:N
ljpart(jj-1)=popweight(jj)*phi(jj)*lam(jj-1)^(-sigmaC/sigmaL)*(phi(jj)/phi(1))^(1/sigmaL);
end
lstst=estst/(popweight(1)*phi(1)+sum(ljpart)); %express l1 using the expression for total effective labor and l2=f(lam,l1)
for jj=2:N
    L2a(jj-1) = lamda(jj-1)^(-sigmaC/sigmaL)*(phi(jj)/phi(1))^(1/sigmaL); %last terms in (11) and f_l
end

tauK0 = taumax;
tolf0 = 1e-7;

if ii==21 | ii==22 | ii==23 | ii==24 | ii==25 | ii==26 | ii==27 | ii==29 | ii==30
    [X,check] = broydn('resid3v',X,tolf0)
    if size(X,1)>1
        X=X';
    end
    sum(resid3v(X).^2)
    for i=1:T;
    time1 = i/T;
    X(i) = time1*(kstst) + (1-time1)*kstart; %k
    end
end

if ii==28 | ii==31 | ii==32
    [X,check] = broydn('resid3v',X,tolf0)
    if size(X,1)>1
        X=X';
    end
    sum(resid3v(X).^2)
end

options = optimset('MaxFunEvals',10000000,'MaxIter',100,'TolFun',1e-15,'TolX',1e-15,'Disp','iter');
[X,check] = fsolve('resid3v',X,options)

if ii==31 | ii==32
for i=1:T;
    time1 = i/T;
    X(i) = time1*(kstst) + (1-time1)*kstart; %k
end
for i=1:T;
    X(2*T+i) = max(5-0.1*i,0); %gamma
end;
options = optimset('MaxFunEvals',10000000,'MaxIter',100,'TolFun',1e-15,'TolX',1e-15,'Disp','iter');
[X,check] = fsolve('resid3v',X,options)
end

if ii==30
[X,check] = broydn('resid3v',X,tolf0)
if size(X,1)>1
    X=X';
end
sum(resid3v(X).^2)
options = optimset('MaxFunEvals',10000000,'MaxIter',100,'TolFun',1e-15,'TolX',1e-15,'Disp','iter');
[X,check] = fsolve('resid3v',X,options)

for i=1:T;
    time1 = i/T;
    X(i) = time1*(kstst) + (1-time1)*kstart; %k
end
for i=1:T;
    X(2*T+i) = max(5-0.1*i,0); %gamma
end;
options = optimset('MaxFunEvals',10000000,'MaxIter',100,'TolFun',1e-15,'TolX',1e-15,'Disp','iter');
[X,check] = fsolve('resid3v',X,options)
end
 
if ii==11 | ii==12 | ii==17 | ii==18 | ii==19
    for i=1:T;
    time1 = i/T;
    X(i) = time1*(kstst) + (1-time1)*kstart; %k
    end
end

if ii<=19
[X,check] = broydn('resid3v',X,tolf0)
if size(X,1)>1
    X=X';
end
sum(resid3v(X).^2)

if ii>=2
for i=1:T;
    time1 = i/T;
    X(i) = time1*(kstst) + (1-time1)*kstart; %k
end
options = optimset('MaxFunEvals',10000000,'MaxIter',100,'TolFun',1e-15,'TolX',1e-15,'Disp','iter');
[X,check] = fsolve('resid3v',X,options)
end
end
% X1=X;
% X=X1;

XX(ii,:)=X;

%for ii=1:length(ALPHall)
%X=XX(ii,:);
k = X(1:T);
l = X(T+1:2*T);
gamma = X(2*T+1:3*T);
Delta1 = X(3*T+1);
Delta2 = X(3*T+2:3*T+2+N-2);
lamda = X(3*T+3+N-2:3*T+3+2*(N-2));
Delta1*kistart(1)+sum(Delta2.*kistart(2:N))
kstst = fzero('findstst',[kstart-0.2 kmax],options)%,0,0)


% save('solN2.mat')
% clc;clear all;close all;
% load('solN2.mat')

rstst = betta^(-1)-1+depr; %r from Euler at the steady state
estst = (rstst/alpha)^(1/(1-alpha))*kstst; % express total effective labor from r=F_k
for jj=2:N
ljpart(jj-1)=popweight(jj)*phi(jj)*lam(jj-1)^(-sigmaC/sigmaL)*(phi(jj)/phi(1))^(1/sigmaL);
end
lstst=estst/(popweight(1)*phi(1)+sum(ljpart)); %express l1 using the expression for total effective labor and l2=f(lam,l1)
for jj=2:N
    L2a(jj-1) = lamda(jj-1)^(-sigmaC/sigmaL)*(phi(jj)/phi(1))^(1/sigmaL); %last terms in (11) and f_l
end
wstst = (1-alpha)*kstst^alpha*estst^(-alpha); %w=F_e
cstst = 1/(popweight(1)+sum(popweight(2:N).*lamda))*(kstst.^alpha*estst.^(1-alpha) - depr*kstst - govexp); %c from resource constraint
mustst = omega*lstst^sigmaL*(1+sum(ALPHA.*lamda.^(-sigmaC).*phi(2:N)'/phi(1).*L2a)+(Delta1+sum(Delta2.*phi(2:N)'/phi(1).*L2a))*(1+sigmaL))/(wstst*(popweight(1)*phi(1)+sum(ljpart)));

klast = [kstart X(1:T-1)]; %k_{t-1}
e=[l' (l'*(lamda.^(-sigmaC/sigmaL).*(phi(2:N)/phi(1))'.^(1/sigmaL)))]*(popweight'.*phi);
e=e';

r = alpha*klast.^(alpha-1).*e.^(1-alpha); %r=F_k
w = (1-alpha)*klast.^(alpha).*e.^(-alpha); %w=F_e
c = 1/(popweight(1)+sum(popweight(2:N).*lamda))*(klast.^alpha.*e.^(1-alpha)+ (1-depr)*klast - k - govexp); %resource constraint

cnext = [c(2:end) cstst];
rnext = [r(2:end) rstst];
lnext = [X(T+2:2*T) lstst];

%check bound on tauk
constrtest = c.^(-sigmaC)./cnext.^(-sigmaC) - betta*(1+ (rnext-depr)*(1-taumax));

cistst = [cstst lamda*cstst]; %consumption of groups in the long run
listst = [lstst, L2(lamda,lstst,phi(2:N))']; %hours worked of groups in the long run
for tt=1:T
ci(tt,:) = [c(tt) lamda*c(tt)]; %consumption of groups during the transition
li(tt,:) = [l(tt), L2(lamda,l(tt),phi(2:N))']; %hours worked of groups during the transition
end

%lifetime utilities
if (sigmaC==1)
    V = sum(betta.^[0:T-1]'*ones(1,N).*(log(ci)- omega*li.^(1+sigmaL)/(1+sigmaL)))...
    + betta^(T)/(1-betta)*(log(cistst) - omega*listst.^(1+sigmaL)/(1+sigmaL));
else
    V = sum(betta.^[0:T-1]'*ones(1,N).*((ci.^(1-sigmaC)-1)/(1-sigmaC) - omega*li.^(1+sigmaL)/(1+sigmaL)))...
    + betta^(T)/(1-betta)*((cistst.^(1-sigmaC)-1)/(1-sigmaC) - omega*listst.^(1+sigmaL)/(1+sigmaL));
end

if sum(V>=Vsq)==2
    PI(ii)=1
else
    PI(ii)=0
end

residall(ii)=sum(resid3v(X).^2)
Vall(ii,:)=V
Vsq

%tax rates
tauKt = 1 - ((c(1:end-1)./c(2:end)).^(-sigmaC)/betta-1)./(r(2:end)-depr); %Euler
tauKt = [tauK0 tauKt];
plot(tauKt)
PI(ii)
tauKt
[swi dur]=min((tauKt-taumax/2).^2);
duration(ii)=dur
lamda_all(ii)=lamda;
Delta1_all(ii)=Delta1;
Delta2_all(ii)=Delta2;
term_all(ii)=1+ALPHA*lamda^(1-sigmaC)+(Delta1+Delta2*lamda)*(1-sigmaC);
kappa=lamda^(-sigmaC/sigmaL)*(phi(2)/phi(1))^(1/sigmaL);
terml_all(ii)=1+ALPHA*kappa^(1+sigmaL)+(Delta1+phi(2)/phi(1)*kappa*Delta2)*(1+sigmaL);
cst_all(ii)=cstst;
gammast_all(ii)=gamma(T);
must_all(ii)=mustst;

%foc for labor
klam=lamda^(-sigmaC/sigmaL)*(phi(2)/phi(1))^(1/sigmaL);
lamderl(ii)=-lamda^(-sigmaC/sigmaL)*(phi(2)/phi(1))^((1+sigmaL)/sigmaL)*betta^T*(1+sigmaL)*lstst^(sigmaL)/(sum(betta.^[0:T-1].*c.^(-sigmaC).*c)+betta^T/(1-betta)*cstst^(-sigmaC)*cstst...
    -sigmaC/sigmaL*lamda^(-(sigmaC+sigmaL)/sigmaL)*(phi(2)/phi(1))^((1+sigmaL)/sigmaL)*(sum(betta.^[0:T-1].*(-omega*l.^(sigmaL).*l))+betta^T/(1-betta)*(-omega)*lstst^sigmaL*lstst));
%foc for consumption
lamderc(ii)=-lamda*betta^T*(1-sigmaC)*cstst^(-sigmaC)/(sum(betta.^[0:T-1].*c.^(-sigmaC).*c)+betta^T/(1-betta)*cstst^(-sigmaC)*cstst...
    -sigmaC/sigmaL*lamda^(-(sigmaC+sigmaL)/sigmaL)*(phi(2)/phi(1))^((1+sigmaL)/sigmaL)*(sum(betta.^[0:T-1].*(-omega*l.^(sigmaL).*l))+betta^T/(1-betta)*(-omega)*lstst^sigmaL*lstst));
%end

save('solN2_dur.mat')

end


clc;clear all;close all;
load('solN2_dur.mat')

Vcap=Vall(:,1)
Vwor=Vall(:,2)

taxesall=zeros(length(ALPHall),T,2);
for ii=1:length(ALPHall)
X=XX(ii,:);

k = X(1:T);
l = X(T+1:2*T);
gamma = X(2*T+1:3*T);
Delta1 = X(3*T+1);
Delta2 = X(3*T+2:3*T+2+N-2);
lamda = X(3*T+3+N-2:3*T+3+2*(N-2));

kstst = fzero('findstst',[kstart-0.2 kmax],options);%,0,0)
rstst = betta^(-1)-1+depr; %r from Euler at the steady state
estst = (rstst/alpha)^(1/(1-alpha))*kstst; % express total effective labor from r=F_k
for jj=2:N
ljpart(jj-1)=popweight(jj)*phi(jj)*lam(jj-1)^(-sigmaC/sigmaL)*(phi(jj)/phi(1))^(1/sigmaL);
end
lstst=estst/(popweight(1)*phi(1)+sum(ljpart)); %express l1 using the expression for total effective labor and l2=f(lam,l1)
for jj=2:N
    L2a(jj-1) = lamda(jj-1)^(-sigmaC/sigmaL)*(phi(jj)/phi(1))^(1/sigmaL); %last terms in (11) and f_l
end
wstst = (1-alpha)*kstst^alpha*estst^(-alpha); %w=F_e
cstst = 1/(popweight(1)+sum(popweight(2:N).*lamda))*(kstst.^alpha*estst.^(1-alpha) - depr*kstst - govexp); %c from resource constraint

klast = [kstart X(1:T-1)]; %k_{t-1}
e=[l' (l'*(lamda.^(-sigmaC/sigmaL).*(phi(2:N)/phi(1))'.^(1/sigmaL)))]*(popweight'.*phi);
e=e';
r = alpha*klast.^(alpha-1).*e.^(1-alpha); %r=F_k
w = (1-alpha)*klast.^(alpha).*e.^(-alpha); %w=F_e
c = 1/(popweight(1)+sum(popweight(2:N).*lamda))*(klast.^alpha.*e.^(1-alpha)+ (1-depr)*klast - k - govexp); %resource constraint

cistst = [cstst lamda*cstst]; %consumption of groups in the long run
listst = [lstst, L2(lamda,lstst,phi(2:N))']; %hours worked of groups in the long run
for tt=1:T
ci(tt,:) = [c(tt) lamda*c(tt)]; %consumption of groups during the transition
li(tt,:) = [l(tt), L2(lamda,l(tt),phi(2:N))']; %hours worked of groups during the transition
end

tauKt = 1 - ((c(1:end-1)./c(2:end)).^(-sigmaC)/betta-1)./(r(2:end)-depr); %Euler
tauKt = [tauK0 tauKt];
tauKtall(ii,:) = tauKt;
[swi dur]=min((tauKt-taumax/2).^2);
duration(ii)=dur;
tauLt = 1 - omega/phi(1)*l.^sigmaL./(c.^(-sigmaC).*w); %consumption-leisure FOC
tauLtall(ii,:) = tauLt;
taxesall(ii,:,:) = [tauKt' tauLt'];
tauLssall(ii) = 1 - omega/phi(1)*lstst^sigmaL/(cstst^(-sigmaC)*wstst);
%tax revenues at each time t
tauKrev = tauKt.*(r-depr).*[kstart k(1:end-1)];
tauLrev = tauLt.*w.*(li*(popweight'.*phi))';
totaltaxrev=tauKrev+tauLrev;
taxrev=[tauKrev' tauLrev'];
deficit(ii,:)=govexp - totaltaxrev;
%PV of taxes
PVtaxes(ii,:) = sum((betta.^[0:T-1].*(c/c(1)).^(-sigmaC))'*ones(1,2).*taxrev)+ betta^T/(1-betta)*(cstst/c(1))^(-sigmaC)*ones(1,2).*taxrev(end,:);
lamall(ii)=lamda;

end
duration;
tauKshare=PVtaxes(:,1)./sum(PVtaxes')';
% lamnor=lamall./(1+lamall);
% psinor=ALPHall./(1+ALPHall);

ALPHall
duration
tauKshare
% lamnor
% psinor

term1=(1-sigmaC)*((1-betta)*Vcap+omega*listst_sq(1)^(1+sigmaL)/(1+sigmaL));
epscap=((term1+1).^(1/(1-sigmaC))/cistst_sq(1)-1)*100
term2=(1-sigmaC)*((1-betta)*Vwor+omega*listst_sq(2)^(1+sigmaL)/(1+sigmaL));
epswor=((term2+1).^(1/(1-sigmaC))/cistst_sq(2)-1)*100

figure
subplot(311)
stairs(epswor,duration,'-k','LineWidth',2);
axis([0 max(epswor) 15 25])
%xlabel('workers` welfare increase (percent)')
%ylabel('duration of transition (years)')
title('Duration of transition (years)')

subplot(312)
tauKshare=PVtaxes(:,1)./sum(PVtaxes')'
plot(epswor,tauKshare,'-k','LineWidth',2);
axis([0 max(epswor) 0.15 0.22])
%xlabel('workers` welfare increase (percent)')
%ylabel('revenue share of capital taxes')
title('Share of capital taxes in goverment revenues')

subplot(313)
alall=ALPHall./(ALPHall+1);
laall=lamall.^sigmaC./(lamall.^sigmaC+1);
plot(epswor,laall,'-k','LineWidth',2);
axis([0 max(epswor) 0.2 0.35])
hold on
plot(epswor,alall,'--k','LineWidth',2);
xlabel("workers' welfare increase (percent)")
legend('\lambda^{\sigma} /(1+\lambda^{\sigma} )', '\psi/(1+\psi)')
title("Workers' relative Pareto weight and consumption share (normalized)")


